package com.hbnrtech.sip.gb28181.handler;

import cn.hutool.core.util.RandomUtil;
import cn.hutool.crypto.digest.MD5;
import com.hbnrtech.sip.gb28181.config.SipConfigProperties;
import com.hbnrtech.sip.gb28181.deivce.DeviceRegistry;
import com.hbnrtech.sip.gb28181.handler.service.SipQueryService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.sip.InvalidArgumentException;
import javax.sip.RequestEvent;
import javax.sip.ServerTransaction;
import javax.sip.SipException;
import javax.sip.address.Address;
import javax.sip.address.AddressFactory;
import javax.sip.address.SipURI;
import javax.sip.header.*;
import javax.sip.message.Request;
import javax.sip.message.Response;
import java.security.NoSuchAlgorithmException;
import java.text.ParseException;
import java.time.Instant;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author liuduo
 * @date 2025-04-21 10:56
 */
@Component
@Slf4j
public class SipRegisterHandler extends SipHandler {
    @Autowired
    private SipConfigProperties sipConfig;
    @Autowired
    private DeviceRegistry deviceRegistry;
    @Autowired
    private HeaderFactory headerFactory;
    @Autowired
    private AddressFactory addressFactory;
    @Autowired
    private SipQueryService sipQueryService;

    // 存储 nonce 以防止重放攻击 (简单实现)
    private final ConcurrentHashMap<String, Instant> nonceMap = new ConcurrentHashMap<>();
    private static final long NONCE_VALIDITY_SECONDS = 300; // nonce 有效期 5 分钟

    public void handleRegister(RequestEvent requestEvent, ServerTransaction serverTransaction) throws SipException, InvalidArgumentException, ParseException {

        Request request = requestEvent.getRequest();

        // 获取必要头信息
        ViaHeader viaHeader = getViaHeader(requestEvent);
        FromHeader fromHeader = getFromHeader(requestEvent);
        ToHeader toHeader = getToHeader(requestEvent); // To == From for REGISTER
        ContactHeader contactHeader = getContactHeader(requestEvent);
        ExpiresHeader expiresHeader = request.getExpires(); // May be null or in Contact
        AuthorizationHeader authHeader = getAuthorizationHeader(requestEvent);
        CallIdHeader callIdHeader = getCallIdHeader(requestEvent);
        CSeqHeader cSeqHeader = getCSeqHeader(requestEvent);
        try {
            if (viaHeader == null || fromHeader == null || toHeader == null || callIdHeader == null || cSeqHeader == null) {
                log.warn("REGISTER request missing mandatory headers. From={}, To={}", fromHeader, toHeader);
                sendResponse(Response.BAD_REQUEST, request, serverTransaction, "Missing mandatory headers");
                return;
            }

            // 提取设备ID (通常在 From/To URI 的 user part)
            String deviceId = ((SipURI) fromHeader.getAddress().getURI()).getUser();
            if (deviceId == null || deviceId.isEmpty()) {
                log.warn("Could not extract deviceId from From header: {}", fromHeader.getAddress());
                sendResponse(Response.BAD_REQUEST, request, serverTransaction, "Invalid From header");
                return;
            }

            // 处理认证
            if (sipConfig.isRequireAuth()) {
                if (authHeader == null) {
                    log.info("Auth required but Authorization header missing for device {}", deviceId);
                    sendUnauthorized(request, serverTransaction);
                    return;
                }
                if (!validateAuth(authHeader, request.getMethod(), deviceId)) {
                    log.warn("Authentication failed for device {}", deviceId);
                    // 认证失败通常也回复 401，让设备重试
                    sendUnauthorized(request, serverTransaction);
                    return;
                }
                log.info("Authentication successful for device {}", deviceId);
            }

            // 处理注册有效期
            int expires = 0; // 默认为注销
            if (expiresHeader != null) {
                expires = expiresHeader.getExpires();
            } else if (contactHeader != null && contactHeader.getExpires() != -1) {
                // Expires也可以在Contact头中
                expires = contactHeader.getExpires();
            } else if (contactHeader != null) {
                // 如果都没有指定，但有Contact头，通常认为是一个长期注册（或按服务器策略）
                // GB28181 规范通常要求显式 Expires
                expires = 3600; // 假设默认 1 小时
                log.warn("REGISTER for {} missing Expires, assuming default {}s", deviceId, expires);
            }

            if (expires < 0) { // 无效值
                sendResponse(Response.BAD_REQUEST, request, serverTransaction, "Invalid Expires value");
                return;
            }

            // 获取设备地址信息
            String remoteIp = viaHeader.getHost();
            int remotePort = viaHeader.getPort();
            String transport = viaHeader.getTransport().toUpperCase(); // UDP or TCP

            // Contact 头中的地址可能更准确，尤其是经过 NAT 的情况
            if (contactHeader != null) {
                SipURI contactUri = (SipURI) contactHeader.getAddress().getURI();
                remoteIp = contactUri.getHost();
                remotePort = contactUri.getPort();
                // transport 可以从 contactUri 的 transport 参数获取
                String contactTransport = contactUri.getTransportParam();
                if (contactTransport != null && !contactTransport.isEmpty()) {
                    transport = contactTransport.toUpperCase();
                }
            }

            // 执行注册或注销
            Instant now = Instant.now();
            if (expires > 0) {
                // 注册或续期
                DeviceRegistry.DeviceInfo deviceInfo = new DeviceRegistry.DeviceInfo(
                        deviceId,
                        remoteIp,
                        remotePort,
                        transport,
                        now,
                        now.plusSeconds(expires)
                );
                deviceRegistry.registerDevice(deviceInfo);
                log.info("Device registered/updated: ID={}, IP={}, Port={}, Transport={}, Expires={}s",
                        deviceId, remoteIp, remotePort, transport, expires);

                // 回复 200 OK
                sendRegisterOk(request, serverTransaction, contactHeader, expires);

                // !!! *** 在这里触发目录查询 *** !!!
                // 最好是异步执行，避免阻塞注册流程太久
                CompletableFuture.runAsync(() -> {
                    try {
                        // 延迟一小段时间再查询，给设备一点处理时间 (可选，根据设备表现调整)
                        Thread.sleep(1000);
                        sipQueryService.queryCatalog(deviceInfo); // 调用查询服务
                    } catch (Exception e) {
                        log.error("Failed to initiate Catalog query for device {}", deviceId, e);
                    }
                });

            } else {
                // 注销
                deviceRegistry.unregisterDevice(deviceId);
                log.info("Device unregistered: ID={}", deviceId);
                // 回复 200 OK (注销成功)
                sendRegisterOk(request, serverTransaction, contactHeader, 0); // Expires=0
            }
        } catch (ParseException | InvalidArgumentException | SipException e) {
            log.error("Error processing REGISTER request", e);
            try {
                // 尝试发送内部服务器错误
                sendResponse(Response.SERVER_INTERNAL_ERROR, request, serverTransaction, "Error processing request");
            } catch (Exception ex) {
                log.error("Failed to send error response", ex);
            }
        }
    }

    // 发送 401 Unauthorized 响应
    private void sendUnauthorized(Request request, ServerTransaction st)
            throws ParseException, SipException, InvalidArgumentException {
        Response response;
        response = getMessageFactory().createResponse(Response.UNAUTHORIZED, request);

        // 创建 WWW-Authenticate 头
        String nonce = generateNonce();
        nonceMap.put(nonce, Instant.now()); // 存储 nonce
        WWWAuthenticateHeader wwwAuthHeader = headerFactory.createWWWAuthenticateHeader("Digest");
        wwwAuthHeader.setParameter("realm", sipConfig.getRealm());
        wwwAuthHeader.setParameter("nonce", nonce);
        wwwAuthHeader.setParameter("algorithm", "MD5");
        wwwAuthHeader.setParameter("qop", "auth"); // 更安全的做法，需要客户端支持

        response.setHeader(wwwAuthHeader);

        // 添加 To, From, Call-ID, CSeq, Via (JAIN-SIP 会自动处理一些)
        ToHeader toHeader = (ToHeader) response.getHeader(ToHeader.NAME);
        if (toHeader.getTag() == null) { // 第一次 401 通常不带 tag
            toHeader.setTag(UUID.randomUUID().toString());
        }

        st.sendResponse(response);
        log.debug("Sent 401 Unauthorized response with nonce: {}", nonce);

    }

    // 发送 200 OK 响应
    private void sendRegisterOk(Request request, ServerTransaction st, ContactHeader contactHeader, int expires)
            throws ParseException, InvalidArgumentException, SipException {
        Response response;
        response = getMessageFactory().createResponse(Response.OK, request);
        // 必须包含 Contact 头，反映当前的注册状态
        if (contactHeader != null) {
            // JAIN-SIP 在创建响应时，通常会自动从请求复制 To, From, Call-ID, CSeq, Via
            // 但 Contact 和 Expires 需要明确设置在响应中以确认状态
            ContactHeader responseContact = (ContactHeader) contactHeader.clone(); // 克隆请求中的 Contact
            responseContact.setExpires(expires); // 设置确认的有效期
            response.setHeader(responseContact);
        } else if (expires > 0) {
            // 如果请求中没有Contact（不规范），但我们接受了注册，可以尝试构建一个
            log.warn("REGISTER OK but request was missing Contact header. Constructing one.");
            Address contactAddress = addressFactory.createAddress(request.getHeader(FromHeader.NAME).toString()); // 用From地址代替
            ContactHeader newContact = headerFactory.createContactHeader(contactAddress);
            newContact.setExpires(expires);
            response.setHeader(newContact);
        }

        // 添加 Expires 头，其值应与 Contact 中的 expires 一致
        ExpiresHeader expiresHeader = headerFactory.createExpiresHeader(expires);
        response.setHeader(expiresHeader);

        // 添加 Date 头
        response.setHeader(headerFactory.createDateHeader(java.util.Calendar.getInstance()));

        // To tag 必须有
        ToHeader toHeader = (ToHeader) response.getHeader(ToHeader.NAME);
        if (toHeader.getTag() == null) {
            toHeader.setTag(UUID.randomUUID().toString()); // 如果请求的 To 没有 tag，需生成
        }

        st.sendResponse(response);
        log.debug("Sent 200 OK response for REGISTER (Expires: {})", expires);
    }


    // 验证摘要认证
    private boolean validateAuth(AuthorizationHeader authHeader, String method, String username) {
        String realm = authHeader.getRealm();
        String nonce = authHeader.getNonce();
        String uri = authHeader.getURI().toString();
        String response = authHeader.getResponse();
        String qop = authHeader.getQop(); // 可选
        int nc = authHeader.getNonceCount(); // 可选
        String cnonce = authHeader.getCNonce();

        if (realm == null || nonce == null || uri == null || response == null) {
            log.warn("Authorization header missing required parameters");
            return false;
        }

        // 验证 Realm 是否匹配
        if (!sipConfig.getRealm().equals(realm)) {
            log.warn("Realm mismatch. Expected: {}, Got: {}", sipConfig.getRealm(), realm);
            return false;
        }

        // 验证 Nonce 是否有效且未过期
        Instant nonceIssueTime = nonceMap.get(nonce);
        if (nonceIssueTime == null) {
            log.warn("Nonce not found or already used: {}", nonce);
            return false; // Nonce 无效或已使用（简单实现中，使用后应移除）
        }
        if (Instant.now().isAfter(nonceIssueTime.plusSeconds(NONCE_VALIDITY_SECONDS))) {
            log.warn("Nonce expired: {}", nonce);
            nonceMap.remove(nonce); // 移除过期的 nonce
            return false;
        }
        // 为防止重放，验证通过后应移除 nonce
        // nonceMap.remove(nonce); // 如果需要严格防止重放

        // 获取设备密码 (极度简化)
        String password = deviceRegistry.getPasswordForDevice(username);
        if (password == null || password.isEmpty()) {
            // 如果是从配置读，这里用 sipConfig.getPassword()
            password = sipConfig.getPassword();
        }
        if (password == null) {
            log.warn("No password configured for device {}", username);
            return false;
        }

        // 计算期望的 response
        try {
            String expectedResponse = calculateResponse(username, realm, password, method, uri, nonce, qop, cnonce, nc);
            log.debug("Auth validation: Username={}, Realm={}, Nonce={}, Method={}, URI={}", username, realm, nonce, method, uri);
            log.debug("Auth validation: Received Response={}, Expected Response={}", response, expectedResponse);
            return response.equals(expectedResponse);
        } catch (NoSuchAlgorithmException e) {
            log.error("MD5 Algorithm not found for Digest Authentication", e);
            return false;
        }
    }

    // 计算 Digest Response
    private String calculateResponse(String username, String realm, String password,
                                     String method, String uri, String nonce,
                                     String qop, String cnonce, int nonceCount) throws NoSuchAlgorithmException {
        // HA1 = MD5(username:realm:password)
        MD5 md5 = MD5.create();
        String ha1 = md5.digestHex(username + ":" + realm + ":" + password);

        // HA2 = MD5(method:uri)
        String ha2 = md5.digestHex(method + ":" + uri);

        // --- Calculate response based on qop ---
        String kd; // Key Digest string
        if (qop != null && (qop.equalsIgnoreCase("auth") || qop.equalsIgnoreCase("auth-int"))) {
            // qop is specified, requires nonceCount and clientNonce
            if (cnonce == null || nonceCount <= 0) {
                throw new IllegalArgumentException("cnonce and positive nonceCount are required when qop is specified.");
            }
            // Format nonceCount as an 8-digit hex string, lowercase (as common)
            String ncString = String.format("%08x", nonceCount);

            // response = MD5(HA1:nonce:nonceCount:clientNonce:qop:HA2)
            kd = ha1 + ":" + nonce + ":" + ncString + ":" + cnonce + ":" + qop + ":" + ha2;

        } else {
            // qop is not specified or not supported, use simplified RFC 2069 style
            // response = MD5(HA1:nonce:HA2)
            kd = ha1 + ":" + nonce + ":" + ha2;
        }

        // Final response is MD5 of the Key Digest string
        return md5.digestHex(kd);
    }

    // 生成唯一的 nonce
    private String generateNonce() {
        // 可以用更复杂的方式保证唯一性和安全性
        return RandomUtil.randomNumbers(16);
    }
}
